{-# OPTIONS -XNamedFieldPuns  #-}

module Sim.Engine (
   Factory,
   Image,
   runProcess,
   activateDetector,
   initImage,
   darkNoiseModeling,
   cosmicTracesModeling,
   obsObjectsModeling,
   bloomingModeling,
   adcModeling,
   mergeImage
 ) where

import  Control.Applicative ( (<$>) )
import Control.Monad.Reader ( runReaderT, ReaderT, local, asks, forM_, liftIO )

import     Data.VectorSpace ( (^/) )
import      Data.Array.Base ( unsafeRead, unsafeWrite )
import        Data.Array.IO ( IOUArray, newArray, getBounds, readArray, writeArray, mapIndices )
import              Data.Ix ( Ix, range, inRange )
import             Data.Map ( fromList, (!) )
import          Text.Printf ( printf )

import              GHC.Arr ( unsafeIndex )
import            GHC.Float ( int2Double )

import       GSL.Random.Gen ( newRNG, mt19937 )
import      GSL.Random.Dist ( getPoisson, getGaussian )

import                 Fits ( Record(..), Keyword(..), Value(..) )

import            Sim.Types
import           Sim.Helper ( unsafeSequenceIO, ergsToPhotons, parEval, mkChan )
import        Sim.CRayTrace ( genCRTraces  )
import        Sim.ObsObject ( toPixObject )
import        Sim.PixObject ( pixelization )

type Image   = IOUArray (Int,Int) Double
type Factory = ReaderT Context IO


-- Запуск основного цикла процесса моделирования
runProcess process c = do

  printCurrentSimOptions c 

  liftIO $ putStrLn "\n---====    Running the simulation   ====---"

  imgs <- sequence [ runReaderT (activateDetector ccd $ pr) c | ccd <- ccdArray $ dev c ]

  liftIO $ putStrLn "\n---====     Simulation completed    ====---\n"

  return (hd,fn,imgs)

  where
    pr = process >>= prepImage

    fn = output $ env c

    hd = [ Record (Keyword "TELESCOP") (VString           "T-170M")   ""
         , Record (Keyword "INSTRUME") (VString (   title $ dev c))   ""
         , Record (Keyword "EXPTIME")  (VDouble ( expTime $ env c))   "Exposure time (sec)"
         , Record (Keyword "CRAYFLUX") (VDouble (  crFlux $ env c))   "Cosmic ray flux"
         ]



-- Активация ПЗС приемника
activateDetector ccd = local (setCurrentCCD ccd)
  where
     setCurrentCCD ccd c = c { ccd }

-- Инициализация изображения
initImage :: Factory Image
initImage = do
  p <- asks $ position.ccd
  w <- asks $    width.ccd
  h <- asks $   height.ccd
  b <- asks $     bias.ccd

  let lr = (p, p+(w,h)-1)

  liftIO $ do
    putStrLn "\nImage initilization..."
    newArray lr b

-- Моделирование теплового шума ПЗС приемника

darkNoiseModeling i = do
  et <- asks $   expTime.env -- время экспозиции
  dn <- asks $ darkNoise.ccd -- значение темнового тока

  liftIO $ do
       putStrLn "DarkNoise modeling..."

       b <- getBounds i
       n <- genDarkNoise (et * dn)

       let ps = zip (range b) n

       forM_ ps $ modifyPixel i (+)

       return i

-- Моделирование следов от воздействия космических частиц

cosmicTracesModeling i = do
  a  <- asks $      area.ccd -- площадь ПЗС приемника
  t  <- asks $ thickness.ccd -- толщина ПЗС приемника
  et <- asks $   expTime.env -- время экспозиции
  fl <- asks $    crFlux.env -- величина потока космических частиц

  liftIO $ do
       putStrLn "CosmicTraces modeling..."
       b <- getBounds i

       traces <- genCRTraces (b,t)

       let n  = round $ a * fl * et
           ps = concatMap (pixelization undefined undefined) $ take n traces

       forM_ ps $ modifyPixel i (+)

       return i

-- Моделирование процесса регистрации наблюдаемых объектов

obsObjectsModeling i = do

  psf' <- asks $       psf.ccd -- функция рассеяния точки (ФРТ)
  adf  <- asks $ disperser.dev -- функция угловой дисперсии
  et   <- asks $   expTime.env -- время экспозиции
  obj  <- asks $      objs.obj -- наблюдаемые объекты

  mc   <- magConverter
  ec   <- sedConverter

  liftIO $ do
       putStrLn "Observation objects modeling..."

       let os = map (toPixObject et mc ec) obj
           ps = concatMap (pixelization psf' adf) os

--       es <- mkChan 2 >>= parEval ps

       forM_ ps $ modifyPixel i (+)

       return i
  where
     magConverter :: Factory (SpType -> Magnitude -> PhotonFlux)
     magConverter = do
          qe       <- asks $  quantumEfficiency.ccd -- квантовая эффективность ПЗС
          sen      <- asks $        sensitivity.dev -- оптическая чувствительность инструмента
          (v, sps) <- asks $ atlasSpectralTypes.obj -- библиотека спектральных типов
          f        <- asks $          magFilter.obj -- фильтр в котором задана звездная величина
          quad     <- asks $               quad.int -- интегратор

          let atlas = fromList $ map proc sps
              proc (st, o) = (st, \m -> k * 10**(-0.4*m) * a / b)
                where
                  a = quad $ v * f
                  b = quad $ o * f
                  k = quad $ ergsToPhotons $ sen * qe * o

          return $ \st -> atlas ! st

     sedConverter :: Factory (SpType -> EnergyDist)
     sedConverter = do
          qe       <- asks $  quantumEfficiency.ccd -- квантовая эффективность ПЗС
          sen      <- asks $        sensitivity.dev -- оптическая чувствительность инструмента
          (_, sps) <- asks $ atlasSpectralTypes.obj -- библиотека спектральных типов
          quad     <- asks $               quad.int -- интегратор


          let atlas = fromList $ map proc sps
              proc (st, o) = (st, f ^/ quad f) 
                where f = ergsToPhotons $ sen * qe * o

          return $ \st -> atlas ! st

-- Моделирование переноса заряда на соседние пиксели при накоплении
-- максимально возможного заряда в пикселе.

bloomingModeling i = do
  l <- asks $ maxLevel.ccd -- максимальное значение в одном пиксиле

  liftIO $ do 
    putStrLn "Blooming effect modeling..."

    b <- getBounds i

    forM_ (range b) $ \p -> do
        c <- readArray i p
        if c > l
           then do
              let n = 0.5 * (c-l)
                  t = \(x,y) -> (x,y+1)
                  b = \(x,y) -> (x,y-1)

              modifyPixel i (-) (p, 2*n)

              transfer i t l n (t p)
              transfer i b l n (b p)
           else return ()

  return i

  where
    -- Функции переноса заряда вверх и вниз соответственно
    transfer img f l n p = do
      b <- getBounds img

      if inRange b p && n /= 0
         then do
           c <- readArray img p

           let dc = if c + n > l then l-c else n

           if c < l
              then do
                modifyPixel img (+) (p, dc)

                transfer img f l (n-dc) (f p)
              else do
                transfer img f l     n  (f p)

         else return()

-- Моделирование процесса аналого-цифрового преобразования (АЦП) ПЗС приемника

adcModeling i = do
  rn <- asks $ readNoise.ccd -- получаем значение шума считывания

  liftIO $ do
       putStrLn "ADC modeling..."

       b <- getBounds i
       n <- genReadNoise rn

       let ps = zip (range b) n

       forM_ ps $ modifyPixel i (+)

       return i

-- Вспомогательные функции

-- Печать основных параметров симуляции
printCurrentSimOptions :: Context -> IO ()
printCurrentSimOptions c = do
  printf "Current simulation options:\n"
  printf "  Exposition Time        : %.2f s\n" ext
  printf "  Cosmic Ray Flux        : %.2f\n" crf

  case psf of
     Gauss fwhm  -> do 
                printf "  Using Gaussian PSF     : True\n"
                printf "  Gaussian FWHM (microns): %s\n" $ show fwhm
     CustomPSF fn -> do
                printf "  Using Custom PSF       : %s\n" $ fn

  case fld of
     RandomStars num  -> do 
                printf "  Random Stars Field     : True\n"
                printf "  Number of Random Stars : %d\n" num

     CustomField fn -> do
                printf "  Using Custom Field     : %s\n" $ fn

 where
   ext = expTime $ env c
   crf = crFlux  $ env c
   psf = psfType $ env c
   fld = fldType $ env c

prepImage i = do

  p <- asks $ ffposition.ccd
  w <- asks $      width.ccd
  h <- asks $     height.ccd

  liftIO $ do
   
    p' <- fst <$> getBounds i

    let dp = p' - p

    mapIndices (p,p+(w,h)-1) (+dp) i

modifyPixel :: Image -> (Double -> Double -> Double) -> ((Int,Int), Double) -> IO ()
modifyPixel img f (i,v) = do
  b <- getBounds img
  case inRange b i of
    False -> return ()
    True  -> unsafeModifyPixel img f (unsafeIndex b i, v)

unsafeModifyPixel :: Image -> (Double -> Double -> Double) -> (Int,Double) -> IO ()
unsafeModifyPixel img f (i,v) =
  do
    c <- unsafeRead img i
    unsafeWrite img i (f c v)

genDarkNoise :: Double -> IO [Double]
genDarkNoise bg = do
     rng <- newRNG mt19937
     unsafeSequenceIO $ repeat $ fmap int2Double $ getPoisson rng bg

genReadNoise :: Double -> IO [Double]
genReadNoise rn = do
     rng <- newRNG mt19937
     unsafeSequenceIO $ repeat $ getGaussian rng rn

mergeImage :: [Image] -> IO Image
mergeImage []     = error "No images to merge!"
mergeImage (x:[]) = return x
mergeImage (x:xs) = do
  x' <- mergeImage xs

  b  <- getBounds x
  b' <- getBounds x'


{--
  print b
  print b'
  print $ unionBounds b b'
--}

  a <- newArray (unionBounds b b') 0

  forM_ (range b') $ \i -> readArray x' i >>= writeArray a i
  forM_ (range b ) $ \i -> readArray x  i >>= writeArray a i

  return a
  where
     unionBounds ((x,y),(x',y')) ((u,v),(u',v')) = ((min x u, min y v), (max x' u', max y' v'))